#include "functions.h"



std::vector<Tins::IP::address_type> srcAddrs; 
std::vector<Tins::IP::address_type> dstAddrs;

unsigned int pduCount[NUM_PDUS]; 
const std::string pduNames[] = { "RAW", "ETHERNET_II", "IEEE802_3", "RADIOTAP", 
  "DOT11", "DOT11_ACK", "DOT11_ASSOC_REQ", "DOT11_ASSOC_RESP", 
  "DOT11_AUTH", "DOT11_BEACON", "DOT11_BLOCK_ACK", "DOT11_BLOCK_ACK_REQ", 
  "DOT11_CF_END", "DOT11_DATA", "DOT11_CONTROL", "DOT11_DEAUTH", 
  "DOT11_DIASSOC", "DOT11_END_CF_ACK", "DOT11_MANAGEMENT", "DOT11_PROBE_REQ", 
  "DOT11_PROBE_RESP", "DOT11_PS_POLL", "DOT11_REASSOC_REQ", "DOT11_REASSOC_RESP", 
  "DOT11_RTS", "DOT11_QOS_DATA", "LLC", "SNAP", 
  "IP", "ARP", "TCP", "UDP", 
  "ICMP", "BOOTP", "DHCP", "EAPOL", 
  "RC4EAPOL", "RSNEAPOL", "DNS", "LOOPBACK", 
  "IPv6", "ICMPv6", "SLL", "DHCPv6", 
  "DOT1Q", "PPPOE", "STP", "PPI", 
  "IPSEC_AH", "IPSEC_ESP", "USER_DEFINED_PDU"
}; 


bool userY()
{
	// Gets input from user
	std::string inputString; 
	std::cin >>  inputString;

	// Compares it to the candidates for "yes" 
	if(inputString.compare("yes") == 0 || inputString.compare("YES") == 0 || inputString.compare("y") == 0){
		return true; 
	}
	else{
		// Anything that does not match the above candidates 
		// is treated like a "no" 
		return false;
	}
}


// Receives a parent PDU and counts that PDU and all child
// PDUs in pduCount
void updateStats(PDU * tempPDU)
{
	pduCount[tempPDU->pdu_type()]++;		// Increments the counter for the parent PDU type

	PDU *innerPDU = tempPDU->release_inner_pdu();	// Gets the child pdu out
	
	// Loops through all children of the parent PDU
	// The fucntion release_inner_pdu returns 0 when there
	// are not childer left to return, hence the conditional 
	// statement in this for-loop
	for(int i=0; innerPDU != 0; i++){
		pduCount[innerPDU->pdu_type()]++; 	// Increment pdu type count

		// If the PDU is of type IP we record the 
		// source and destination addresses if they
		// are new
		if(innerPDU->pdu_type() == IP_NUM){
			// Update source IPs
			if(std::find(srcAddrs.begin(), srcAddrs.end(), ((IP*)(innerPDU))->src_addr()) == srcAddrs.end()){
				// Was not found, add it to table
				srcAddrs.push_back(((IP*)(innerPDU))->src_addr()); 
			}
			// Update destination IPs
			if(std::find(dstAddrs.begin(), dstAddrs.end(), ((IP*)(innerPDU))->dst_addr()) == dstAddrs.end()){
				// Was not found, add it to table
				dstAddrs.push_back(((IP*)(innerPDU))->dst_addr()); 
			}
		}
		// Gets child PDU if it exists
		innerPDU = innerPDU->release_inner_pdu(); 
	}
}

// Prints formatted statistics based on the
// information gathered during the sniffing phase
void showStats()
{
	std::cout << std::endl << "Packet Statistics: " << std::endl; 

	// The number of hits for the packet type with most hits
	// must be found in order to make the histogram
	int max = 0; 
	for(int i=0; i<NUM_PDUS; i++){
		if(pduCount[i] != 0){	// Prints protocol name and the number of hits if > 0 
			std::cout << std::setw(10) << std::left << pduNames[i] << "\t" << pduCount[i] << std::endl; 
		}
		if(pduCount[i] > max)
			max = pduCount[i];
	}

	std::cout << std::endl << "Histogram: " << std::endl; 	
	// Prints a histogram with the highest counting packet earning MAX_STARS stars
	for(int i=0; i<NUM_PDUS; i++){
		if(pduCount[i] != 0){
			std::cout << std::setw(10) << std::left<< pduNames[i] << "\t";
			for(int c=0; c<float(pduCount[i])/max*MAX_STARS; c++)
				std::cout << "*"; 
			std::cout << std::endl;  
		}
	}

	
	// Prints the whole srcAddrs vector, which holds 
	// unique source addresses
	std::cout << std::endl << "Unique Source Addresses: " << std::endl; 
	for(int i=0; i<srcAddrs.size(); i++){
		std::cout << srcAddrs[i] << std::endl; 
	}

	// Prints the whole dstAddrs vector, which holds
	// unique destination addresses
	std::cout << std::endl << "Unique Destination Addresses: " << std::endl; 
	for(int i=0; i<dstAddrs.size(); i++){
		std::cout << dstAddrs[i] << std::endl; 
	}

}

// Recursively dissects PDUs and outputs info
// to the commandline
void pduDissector(PDU * tempPDU, int packetNum)
{
	// packetSummary is a string that describes the whole pdu
	// from parent to the last child. It is built as the parent 
	// pdu is traversed
	std::string packetSummary = "[ "; 

	// Outputs the packet number and its type
	std::cout << "Packet#: " << packetNum << "\tType: " << pduNames[tempPDU->pdu_type()] << std::endl;
	packetSummary += pduNames[tempPDU->pdu_type()]; 

	// Gets the child PDU
	PDU *innerPDU = tempPDU->release_inner_pdu();

	// Iterates through all of the children
	// and displays them as "nested packets"
	for(int i=0; innerPDU != 0; i++){
		for(int c=0; c<=i+1; c++)
			std::cout << "----------"; 
		std::cout<< ">  Nested Packet: " << pduNames[innerPDU->pdu_type()] << " ";
		
		// If the PDU is an IP packet we retrieve the 
		// source and destination data from it
		if(innerPDU->pdu_type() == IP_NUM){
			std::cout << "\tsrc: " << ((IP*)(innerPDU))->src_addr() << "\tdest: " << ((IP*)(innerPDU))->dst_addr() << std::endl; 
		}
		else{
			std::cout << std::endl; 
		}

		// Formatting for the packetSummary string
		packetSummary += " | ";  
		packetSummary += pduNames[innerPDU->pdu_type()]; 
		
		// Gets the next child if it exists
		innerPDU = innerPDU->release_inner_pdu();  
	}
	packetSummary += " ]"; 
	std::cout << packetSummary << std::endl << std::endl; 
}
